# %%
from collections import defaultdict
from datetime import datetime, timezone
from itertools import batched
import time
from loguru import logger
from pydantic import BaseModel, Field
from pymongo import InsertOne, UpdateOne
from tqdm import tqdm
from conn import homologous_clip, mysql_engine, clip_col, clip_annotated_record
from sqlalchemy import text


# %%
START_DATE = "2025-05-19 00:00:00"
START_TS = int(datetime.strptime(START_DATE, "%Y-%m-%d %H:%M:%S").timestamp() * 1000)
print(START_DATE, START_TS)

# %%
# 查找和更新送标记录
BATCH = 1000
max_id = int(1e9)
clip_map = defaultdict(set)
with mysql_engine.connect() as conn:
    break_flag = False
    while True:
        sql = text(
            f"""
            select id, clip_id, annotate_type, ctime
            from send_label_record_tab
            where id < :max_id
            order by id desc
            limit {BATCH}
            """
        ).bindparams(max_id=max_id)

        print(f"max_id: {max_id}")
        rows = conn.execute(sql).fetchall()
        max_id = rows[-1][0]
        for id, clip_id, annotate_type, ctime in rows:
            if ctime < START_TS:
                break_flag = True
                break
            clip_map[clip_id].add(annotate_type)

        if break_flag:
            break

# %%
clip_ids = list(clip_map.keys())
print(len(clip_ids))
for batch_clip_ids in tqdm(batched(clip_ids, BATCH), total=len(clip_ids) // BATCH):
    ops = []
    exists_clip_ids = []
    for clip in homologous_clip.find(
        {"clip_id": {"$in": batch_clip_ids}},
        {"clip_id": 1, "curated_annotate_type_list": 1},
    ):
        clip_id, curated_annotate_type_list = (
            clip["clip_id"],
            clip["curated_annotate_type_list"],
        )
        new_curated_annotate_type_list = sorted(
            list(clip_map[clip_id] | set(curated_annotate_type_list))
        )
        exists_clip_ids.append(clip_id)
        ops.append(
            UpdateOne(
                {"_id": clip_id},
                {
                    "$set": {
                        "curated_annotate_type_list": new_curated_annotate_type_list
                    }
                },
            )
        )
    new_clip_ids = list(set(batch_clip_ids) - set(exists_clip_ids))
    if new_clip_ids:
        for clip in clip_col.find(
            {"_id": {"$in": new_clip_ids}},
            {"name": 1, "tags": 1, "collect_time": 1},
        ):
            clip_id = clip.pop("_id")
            ops.append(
                InsertOne(
                    {
                        "clip_id": clip_id,
                        "curated_annotate_type_list": sorted(list(clip_map[clip_id])),
                        "created_at": datetime.now(tz=timezone.utc),
                        **clip,
                    }
                )
            )

    homologous_clip.bulk_write(ops, ordered=False)


# %%
# 查询最近产生的回标 clip
ops = []
clip_ids = []
BATCH = 1000
for clip in clip_annotated_record.find(
    {"ctime": {"$gt": START_TS}},
    {"state": 1, "clip_id": 1, "annotate_type": 1},
):
    clip_ids.append(clip["clip_id"])
    ops.append(
        UpdateOne(
            {"clip_id": clip["clip_id"]},
            {"$addToSet": {"labeled_annotate_type_list": clip["annotate_type"]}},
        )
    )

    if clip["state"] == 2:
        ops.append(
            UpdateOne(
                {"clip_id": clip["clip_id"]},
                {"$addToSet": {"confirmed_annotate_type_list": clip["annotate_type"]}},
            )
        )

    if len(clip_ids) == BATCH:
        print(f"{len(ops) = }")
        homologous_clip.bulk_write(ops)
        ops = []
        clip_ids = []

if ops:
    homologous_clip.bulk_write(ops)
